package org.example.combine;

import org.apache.flink.api.java.tuple.Tuple;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.streaming.api.datastream.ConnectedStreams;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.CoProcessFunction;
import org.apache.flink.util.Collector;
import org.w3c.dom.stylesheets.LinkStyle;

import java.util.*;

/**
 * 连接两条流，实现类似JOIN
 */
public class ConnectKeybyDemo {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        DataStreamSource<Tuple2<Integer, String>> source1 = env.fromElements(
                Tuple2.of(1, "a1"),
                Tuple2.of(1, "a2"),
                Tuple2.of(2, "b"),
                Tuple2.of(3, "c"));
        DataStreamSource<Tuple3<Integer, String, Integer>> source2 = env.fromElements(
                Tuple3.of(1, "aa1", 1),
                Tuple3.of(1, "aa2", 2),
                Tuple3.of(2, "bb", 1),
                Tuple3.of(3, "cc", 3));

        ConnectedStreams<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>> connect = source1.connect(source2);
        ConnectedStreams<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>> keyByStream = connect.keyBy(key1 -> key1.f0, key2 -> key2.f0);
        //两条流不一定谁先来
        //每条流有数据来就先存变量中
        //需要去对方的变量中查找匹配的数据
        SingleOutputStreamOperator<String> process = keyByStream.process(new CoProcessFunction<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>, String>() {
            HashMap<Integer, List<Tuple2<Integer, String>>> map1 = new HashMap<>();
            HashMap<Integer, List<Tuple3<Integer, String, Integer>>> map2 = new HashMap<>();

            @Override
            public void processElement1(Tuple2<Integer, String> value, CoProcessFunction<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>, String>.Context ctx, Collector<String> out) throws Exception {
                map1.putIfAbsent(value.f0, new ArrayList<>());
                map1.get(value.f0).add(value);
                List<Tuple3<Integer, String, Integer>> s2 = map2.get(value.f0);
                if (s2 != null) {
                    for (Tuple3<Integer, String, Integer> tuple3 : s2) {
                        out.collect("s1:" + value + "===>s2:" + tuple3);
                    }
                }
            }

            @Override
            public void processElement2(Tuple3<Integer, String, Integer> value, CoProcessFunction<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>, String>.Context ctx, Collector<String> out) throws Exception {
                map2.putIfAbsent(value.f0, new ArrayList<>());
                map2.get(value.f0).add(value);
                List<Tuple2<Integer, String>> s1 = map1.get(value.f0);
                if (s1 != null) {
                    for (Tuple2<Integer, String> tuple2 : s1) {
                        out.collect("s2:" + value + "===>s1:" + tuple2);
                    }
                }
            }
        });
        process.print();

        env.execute("ConnectKeybyDemo");
    }
}
